package com.example.langchanin4jdemo1.controller;

import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

public class EmbeddingDemo3 {
    public static void main(String[] args) {
        EmbeddingModel embeddingModel = QwenEmbeddingModel
                .builder()
                .apiKey("sk-875dd6ef14244431acdc7ccb974f5bfe")
                .modelName("text-embedding-v2")
                .build();

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536) //维度，需要与计算结果保持⼀致。如果使⽤其他的模型，可能会有不同的结果。
                .build();

        //预设⼏个指示，⽣成向量
        TextSegment textSegment1 = TextSegment.textSegment("客服电话是400-8558558");
        TextSegment textSegment2 = TextSegment.textSegment("客服⼯作时间是周⼀到周五");
        TextSegment textSegment3 = TextSegment.textSegment("客服投诉电话是400-8668668");
        Response<Embedding> embed1 = embeddingModel.embed(textSegment1);
        Response<dev.langchain4j.data.embedding.Embedding> embed2 = embeddingModel.embed(textSegment2);
        Response<dev.langchain4j.data.embedding.Embedding> embed3 = embeddingModel.embed(textSegment3);
        // 存储向量
        embeddingStore.add(embed1.content(), textSegment1);
        embeddingStore.add(embed2.content(), textSegment2);
        embeddingStore.add(embed3.content(), textSegment3);
        // 预设⼀个问题，⽣成向量
        Response<dev.langchain4j.data.embedding.Embedding> embed = embeddingModel.embed("客服电话多少");

        // 查询
        List<EmbeddingMatch<TextSegment>> result = embeddingStore.findRelevant(embed.content(), 5);
        for (EmbeddingMatch<TextSegment> embeddingMatch : result) {
            System.out.println(embeddingMatch.embedded().text() + ",分数为：" +
                    embeddingMatch.score());
        }
    }
}
